# coding=utf-8
# PLIHandle
# Description:
# Created by alvin on 2021/12/14.
# Copyright © 2021 YZZH. All rights reserved.


import os, cv2
from os.path import *
from tqdm import tqdm

provinces = ['皖', '沪', '津', '渝', '冀', '晋', '蒙', '辽', '吉', '黑', '苏', '浙', '京', '闽', '赣', '鲁', '豫', '鄂',
             '湘', '粤', '桂', '琼', '川', '贵', '云', '藏', '陕', '甘', '青', '宁', '新', '警', '学', 'O']

alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W',
             'X', 'Y', 'Z', 'O']

ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
       'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O']


def STRcoordinates_deal(STRcoordinates, i):
    # STRcoordinates = STRcoordinates.split('_')
    result = STRcoordinates[i]
    x, y = int(result.split('&')[0]), int(result.split('&')[1])
    result = x, y
    return result


def analyse_file(filename):
    STRinformation = filename.split('-')
    STRarea_ratio = STRinformation[0]
    STRangle = STRinformation[1]
    STRcoordinates = STRinformation[3]  # 386&473_177&454_154&383_363&402 右下、左下、左上、右下
    STRchars = STRinformation[4]

    # 坐标处理
    STRcoordinates = STRcoordinates.split('_')
    right_bottom = STRcoordinates_deal(STRcoordinates, 0)
    left_top = STRcoordinates_deal(STRcoordinates, 1)
    left_bottom = STRcoordinates_deal(STRcoordinates, 2)
    right_top = STRcoordinates_deal(STRcoordinates, 3)
    coordinates = [left_top, right_top, left_bottom, right_bottom]

    # 车牌字符处理
    LPchars = []
    STRchars = STRchars.split('_')
    LPchars.append(provinces[int(STRchars[0])])
    LPchars.append(alphabets[int(STRchars[1])])
    for i in range(2, 7):
        LPchars.append(ads[int(STRchars[i])])

    return coordinates, LPchars


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def crop_image(image, scope):
    '''将原图片按照scope进行切割，image为被截取的图片，
    scope为需要截取的轮廓的范围，形式为[x1,x2,y1,y2],(x1,y1)为包裹轮廓的直矩形的左上点，（x2,y2）为右下点'''
    x1, x2, y1, y2 = scope
    img = image[y1:y2, x1:x2]
    # cv2.imwrite(path+name,img)
    return img


def ccpd_image_handle(start_index, startSize, totalSize, input_image_path, output_image_path, size):
    output_plate_path = join(output_image_path, 'plates')
    output_cars_path = join(output_image_path, 'images')
    output_labels_path = join(output_image_path, 'labels')
    for path in (output_plate_path, output_cars_path, output_labels_path):
        if not exists(path):
            os.makedirs(path)
    file_paths = list(filter(lambda x: x.lower().endswith('.jpg'), os.listdir(input_image_path)))
    file_paths = file_paths[startSize:min(len(file_paths), startSize + totalSize)]
    for path in tqdm(file_paths):
        img = cv2.imread(join(input_image_path, path))
        result = splitext(path)
        filename = result[0]
        ext = result[1]
        coordinates, lpchars = analyse_file(filename)
        plate = "".join(lpchars)
        imagename = str(start_index).zfill(9)

        # 把图片重命名后保存到对应路径下
        imagefilename = join(output_cars_path, imagename + ext)
        cv2.imwrite(imagefilename, img)

        # 把图片对应的标签保存到对应路径下
        x_list = list()
        y_list = list()
        for x, y in coordinates:
            x_list.append(x)
            y_list.append(y)
        x_min = min(x_list)
        x_max = max(x_list)
        y_min = min(y_list)
        y_max = max(y_list)

        b1, b2, b3, b4 = x_min, x_max, y_min, y_max
        h, w, _ = img.shape
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = convert((w, h), (b1, b2, b3, b4))
        with open(f'{output_labels_path}/{imagename}.txt', 'w') as f:
            f.write("0" + " " + " ".join([str(a) for a in b]) + '\n')

        # 截取图片中车牌号部分
        im = crop_image(img, (b1, b2, b3, b4))
        im = cv2.resize(im, size)
        cv2.imwrite(f'{output_plate_path}/{plate}.jpg', im)  # 保存截图文件的路径

        start_index += 1


if __name__ == "__main__":
    image_path = "/Volumes/Application/{}/ccpd_tilt"
    input_image_path = image_path.format("CCPD2019")
    output_image_path = image_path.format("CCPD2019Handle")

    ccpd_image_handle(start_index=0,
                      startSize=0,
                      totalSize=12800,
                      input_image_path=input_image_path,
                      output_image_path=output_image_path,
                      size=(164, 48)
                      )
